package com.lx.plugin

import com.android.build.api.transform.Format
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformInvocation
import com.android.build.gradle.internal.pipeline.TransformManager
import com.android.utils.FileUtils
import org.apache.commons.compress.utils.IOUtils
import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import java.io.FileOutputStream
import java.nio.file.attribute.FileTime
import java.util.Enumeration
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream
import java.util.zip.CRC32
import java.util.zip.ZipEntry


/**
 * Created by lixiong on 2024/4/28.
 */
class MethodTimePlugin : Transform() {

    private val fileTime = FileTime.fromMillis(0)

    override fun getName(): String {
        return "MethodTimePlugin"
    }

    /**
     * 用于指明Transform的输入类型
     */
    override fun getInputTypes(): MutableSet<QualifiedContent.ContentType> {
        return TransformManager.CONTENT_CLASS
    }

    /**
     * 用于指明Transform的作用域
     */
    override fun getScopes(): MutableSet<in QualifiedContent.Scope> {
        return TransformManager.SCOPE_FULL_PROJECT
    }

    /**
     * 指明该Transform是否支持增量编译
     */
    override fun isIncremental(): Boolean {
        return true
    }

    override fun transform(transformInvocation: TransformInvocation?) {
        super.transform(transformInvocation)
        val inputs = transformInvocation?.inputs
        val outputProvider = transformInvocation?.outputProvider
        inputs?.forEach { transformInput ->
            // 遍历项目目录
            transformInput.directoryInputs.forEach { directoryInput ->
                if (directoryInput.file.isDirectory) {
                    FileUtils.getAllFiles(directoryInput.file).forEach { file ->
                        val name = file.name
                        // 过滤class文件, 排除R.class, BuildConfig.class
                        if (name.endsWith(".class") && !name.startsWith("R\$") &&
                            name != "R.class" && name != "BuildConfig.class"
                        ) {
                            // 找到需要的class文件，进行插桩
                            val path = file.absolutePath
                            val cr = ClassReader(file.readBytes())
                            val cw = ClassWriter(cr, ClassWriter.COMPUTE_MAXS)
                            val visitor = MethodTimeClassVisitor(cw)
                            cr.accept(visitor, ClassReader.EXPAND_FRAMES)

                            val bytes = cw.toByteArray()
                            var fos: FileOutputStream? = null
                            try {
                                fos = FileOutputStream(path)
                                fos.write(bytes)
                            } catch (e: Exception) {
                                e.printStackTrace()
                            } finally {
                                runCatching { fos?.close() }
                            }
                        }
                    }
                }
                val dest = outputProvider?.getContentLocation(
                    directoryInput.name,
                    directoryInput.contentTypes,
                    directoryInput.scopes,
                    Format.DIRECTORY
                )
                FileUtils.copyDirectoryToDirectory(directoryInput.file, dest)
            }

            // 遍历jar包
            transformInput.jarInputs.forEach { jarInput ->
                val dest = outputProvider?.getContentLocation(
                    jarInput.name,
                    jarInput.contentTypes,
                    jarInput.scopes,
                    Format.JAR
                )
                if (dest != null) {
                    FileUtils.mkdirs(dest.parentFile)
                    // 只对 test.jar 进行插桩
                    if (jarInput.file.name.endsWith("test.jar")) {
                        var jos: JarOutputStream? = null
                        try {
                            val jarFile = JarFile(jarInput.file)
                            jos = JarOutputStream(FileOutputStream(dest))
                            val entries: Enumeration<JarEntry> = jarFile.entries()
                            while (entries.hasMoreElements()) {
                                val entry: JarEntry = entries.nextElement()
                                val name: String = entry.name
                                val outEntry = JarEntry(name)
                                val inputStream = jarFile.getInputStream(entry)
                                // 过滤class文件, 排除R.class, BuildConfig.class
                                val newEntryContent = if (name.endsWith(".class") && !name.startsWith("R\$") &&
                                    name != "R.class" && name != "BuildConfig.class"
                                ) {
                                    // 找到需要的class文件，进行插桩
                                    val cr = ClassReader(inputStream)
                                    val cw = ClassWriter(cr, ClassWriter.COMPUTE_MAXS)
                                    val visitor = MethodTimeClassVisitor(cw)
                                    cr.accept(visitor, ClassReader.EXPAND_FRAMES)
                                    cw.toByteArray()
                                } else {
                                    IOUtils.toByteArray(inputStream)
                                }
                                // 将处理后的类文件写入 JAR 包
                                val crc32 = CRC32()
                                crc32.update(newEntryContent)
                                outEntry.crc = crc32.value
                                outEntry.method = ZipEntry.STORED
                                outEntry.size = newEntryContent.size.toLong()
                                outEntry.compressedSize = newEntryContent.size.toLong()
                                outEntry.setLastAccessTime(fileTime)
                                outEntry.setLastModifiedTime(fileTime)
                                outEntry.setCreationTime(fileTime)
                                jos.putNextEntry(outEntry)
                                jos.write(newEntryContent)
                                jos.closeEntry()
                            }
                        } catch (e: Exception) {
                            e.printStackTrace()
                        } finally {
                            runCatching {
                                jos?.flush()
                                jos?.close()
                            }
                        }
                    } else {
                        FileUtils.copyFile(jarInput.file, dest)
                    }
                }
            }
        }
    }
}